package encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
)

func Encrypt(data []byte, key []byte, iv []byte) ([]byte, error) {
	aesBlockEncrypt, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println(err.Error())
		return nil, err
	}

	content := pKCS5Padding(data, aesBlockEncrypt.BlockSize())
	cipherBytes := make([]byte, len(content))
	aesEncrypt := cipher.NewCBCEncrypter(aesBlockEncrypt, iv)
	aesEncrypt.CryptBlocks(cipherBytes, content)
	result := []byte(base64.StdEncoding.EncodeToString(cipherBytes))
	return result, nil
}

func Decrypt(data string, key []byte, iv []byte) (result []byte, err error) {
	// 这里会报数组越界异常，所以  recover 捕获一下
	defer func() {
		if re := recover(); re != nil {
			switch ty := re.(type) {
			case string:
				err = errors.New(ty)
			default:
				err = fmt.Errorf("%v", ty)
			}
		}
	}()
	crytedByte, _ := base64.StdEncoding.DecodeString(data)
	crytedByte1 := []byte(crytedByte)
	decrypted := make([]byte, len(crytedByte))
	var aesBlockDecrypt cipher.Block
	aesBlockDecrypt, err = aes.NewCipher(key)
	if err != nil {
		fmt.Println(err.Error())
		return nil, err
	}
	aesDecrypt := cipher.NewCBCDecrypter(aesBlockDecrypt, iv)
	aesDecrypt.CryptBlocks(decrypted, crytedByte1)
	return pKCS5Trimming(decrypted), nil
}

func pKCS5Padding(cipherText []byte, blockSize int) []byte {
	padding := blockSize - len(cipherText)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(cipherText, padText...)
}

func pKCS5Trimming(encrypt []byte) []byte {
	padding := encrypt[len(encrypt)-1]
	return encrypt[:len(encrypt)-int(padding)]
}

func EncryptToUrlBase64(data []byte, key []byte, iv []byte) ([]byte, error) {
	aesBlockEncrypt, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println(err.Error())
		return nil, err
	}

	content := pKCS5Padding(data, aesBlockEncrypt.BlockSize())
	cipherBytes := make([]byte, len(content))
	aesEncrypt := cipher.NewCBCEncrypter(aesBlockEncrypt, iv)
	aesEncrypt.CryptBlocks(cipherBytes, content)
	result := []byte(base64.RawURLEncoding.EncodeToString(cipherBytes))
	return result, nil
}

// FromBase64UrlDecrypt 不带填充符，适用于url的base64编码，替换了+和/
func FromBase64UrlDecrypt(data string, key []byte, iv []byte) (result []byte, err error) {
	crytedByte, _ := base64.RawURLEncoding.DecodeString(data)
	crytedByte1 := []byte(crytedByte)
	decrypted := make([]byte, len(crytedByte))
	var aesBlockDecrypt cipher.Block
	aesBlockDecrypt, err = aes.NewCipher(key)
	if err != nil {
		fmt.Println(err.Error())
		return nil, err
	}
	aesDecrypt := cipher.NewCBCDecrypter(aesBlockDecrypt, iv)
	aesDecrypt.CryptBlocks(decrypted, crytedByte1)
	return pKCS5Trimming(decrypted), nil
}

func Md5(key any) (string, error) {
	hasher := md5.New()
	var data []byte
	switch v := key.(type) {
	case string:
		data = []byte(v)
	case []byte:
		data = v
	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64:
		data = []byte(fmt.Sprintf("%v", v))
	default:
		var err error
		data, err = json.Marshal(v)
		if err != nil {
			return "", fmt.Errorf("failed to marshal key to JSON: %w", err)
		}
	}
	hasher.Write(data)
	return hex.EncodeToString(hasher.Sum(nil)), nil
}
